#!/usr/bin/env python3
"""
账号池管理器 - 第1.5阶段核心组件
负责管理多个账号的轮换使用，实现IP-账号严格绑定，防止关联封禁
"""

import logging
import time
import yaml
from typing import Dict, List, Optional, Any
from pathlib import Path
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)

class AccountPoolManager:
    """轻量级账号池管理器 - 第1.5阶段"""
    
    def __init__(self):
        """初始化账号池管理器"""
        logger.info("[ACCOUNT_POOL] 初始化账号池管理器...")
        
        try:
            self.accounts = self.load_accounts_config()
            self.account_status = {}  # 账号状态跟踪
            self.usage_stats = {}     # 使用统计
            self.network_pools = {    # 网络环境账号池
                'company': [],        # 公司网络专用账号
                'home': []           # 家庭网络专用账号
            }
            
            # 初始化账号池
            self.init_account_pools()
            
            # 加载安全配置
            self.security_config = self.load_security_config()
            
            logger.info(f"[ACCOUNT_POOL] ✅ 初始化完成，共加载 {len(self.accounts)} 个账号")
            logger.info(f"[ACCOUNT_POOL] 公司账号: {len(self.network_pools['company'])} 个")
            logger.info(f"[ACCOUNT_POOL] 家庭账号: {len(self.network_pools['home'])} 个")
            
        except Exception as e:
            logger.error(f"[ACCOUNT_POOL] ❌ 初始化失败: {e}")
            raise
    
    def load_accounts_config(self) -> List[Dict]:
        """加载账号配置"""
        config_file = Path(__file__).parent.parent / "config" / "accounts.yaml"
        
        try:
            if not config_file.exists():
                logger.warning(f"[ACCOUNT_POOL] 配置文件不存在: {config_file}")
                return self.create_default_config()
            
            with open(config_file, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
                return config.get('accounts', [])
                
        except Exception as e:
            logger.error(f"[ACCOUNT_POOL] 配置文件加载失败: {e}")
            return self.create_default_config()
    
    def create_default_config(self) -> List[Dict]:
        """创建默认配置（用于测试）"""
        logger.info("[ACCOUNT_POOL] 创建默认测试配置")
        
        return [
            {
                'id': 'cookie1',
                'username': 'main_cookie_1',
                'cookie_file': 'config/cookies/cookie1.txt',
                'network_binding': 'main',
                'time_slots': ['morning', 'afternoon', 'evening'],
                'daily_limit': 800,
                'priority': 1
            },
            {
                'id': 'cookie2',
                'username': 'main_cookie_2',
                'cookie_file': 'config/cookies/cookie2.txt',
                'network_binding': 'main',
                'time_slots': ['morning', 'afternoon', 'evening'],
                'daily_limit': 800,
                'priority': 1
            },
            {
                'id': 'main_cookie',
                'username': 'main_cookie_fallback',
                'cookie_file': 'config/cookie.txt',
                'network_binding': 'main',
                'time_slots': ['morning', 'afternoon', 'evening'],
                'daily_limit': 800,
                'priority': 2
            }
        ]

    def get_default_accounts(self) -> List[Dict]:
        """获取默认账号配置（兼容性方法）"""
        return self.create_default_config()

    def load_security_config(self) -> Dict:
        """加载安全配置"""
        try:
            config_file = Path(__file__).parent.parent / "config" / "accounts.yaml"
            if config_file.exists():
                with open(config_file, 'r', encoding='utf-8') as f:
                    config = yaml.safe_load(f)
                    return config.get('security', {})
        except Exception as e:
            logger.warning(f"[ACCOUNT_POOL] 安全配置加载失败: {e}")
        
        # 默认安全配置
        return {
            'max_daily_crawl': 2000,
            'account_cooling_time': 3600,  # 1小时
            'error_cooling_multiplier': 2,
            'max_consecutive_errors': 3,
            'session_timeout': 7200  # 2小时
        }
    
    def init_account_pools(self):
        """初始化账号池分配"""
        for account in self.accounts:
            network_type = account['network_binding']
            
            # 分配到对应网络池
            if network_type in self.network_pools:
                self.network_pools[network_type].append(account)
            else:
                logger.warning(f"[ACCOUNT_POOL] 未知网络类型: {network_type}")
                continue
            
            # 初始化账号状态
            account_id = account['id']
            self.account_status[account_id] = {
                'status': 'available',  # available, in_use, cooling, blocked
                'last_used': 0,
                'usage_count_today': 0,
                'consecutive_errors': 0,
                'total_data_crawled': 0,
                'last_success_time': 0,
                'cooling_until': 0,
                'created_time': time.time()
            }
            
            # 初始化使用统计
            self.usage_stats[account_id] = {
                'total_sessions': 0,
                'successful_sessions': 0,
                'total_data': 0,
                'last_7_days': [],
                'success_rate': 1.0
            }
    
    def get_available_account(self, network_type: str) -> Optional[Dict]:
        """获取指定网络环境的可用账号"""
        logger.debug(f"[ACCOUNT_POOL] 请求 {network_type} 网络的可用账号")
        
        if network_type not in self.network_pools:
            logger.error(f"[ACCOUNT_POOL] 不支持的网络类型: {network_type}")
            return None
        
        available_accounts = []
        
        # 筛选可用账号
        for account in self.network_pools[network_type]:
            account_id = account['id']
            if self.is_account_available(account_id):
                available_accounts.append(account)
        
        if not available_accounts:
            logger.warning(f"[ACCOUNT_POOL] {network_type} 网络无可用账号")
            self.log_pool_status(network_type)
            return None
        
        # 选择最优账号（使用次数最少的）
        selected = min(available_accounts, 
                      key=lambda a: self.account_status[a['id']]['usage_count_today'])
        
        # 标记为使用中
        account_id = selected['id']
        self.account_status[account_id]['status'] = 'in_use'
        self.account_status[account_id]['last_used'] = time.time()
        
        logger.info(f"[ACCOUNT_POOL] ✅ 分配账号: {account_id} ({network_type})")
        return selected
    
    def is_account_available(self, account_id: str) -> bool:
        """检查账号是否可用"""
        if account_id not in self.account_status:
            logger.warning(f"[ACCOUNT_POOL] 账号状态不存在: {account_id}")
            return False
        
        status = self.account_status[account_id]
        current_time = time.time()
        
        # 检查账号状态
        if status['status'] in ['blocked']:
            return False
        
        # 检查冷却时间
        if status['status'] == 'cooling' and current_time < status['cooling_until']:
            return False
        
        # 检查今日使用量限制
        daily_limit = self.get_account_daily_limit(account_id)
        if status['usage_count_today'] >= daily_limit:
            logger.debug(f"[ACCOUNT_POOL] 账号 {account_id} 已达今日限额: {status['usage_count_today']}/{daily_limit}")
            return False
        
        # 检查连续错误次数
        max_errors = self.security_config.get('max_consecutive_errors', 3)
        if status['consecutive_errors'] >= max_errors:
            logger.debug(f"[ACCOUNT_POOL] 账号 {account_id} 连续错误过多: {status['consecutive_errors']}")
            return False
        
        # 检查最小冷却间隔
        min_cooling = self.security_config.get('account_cooling_time', 3600)
        if current_time - status['last_used'] < min_cooling:
            return False
        
        return True
    
    def get_account_daily_limit(self, account_id: str) -> int:
        """获取账号每日限额"""
        for account in self.accounts:
            if account['id'] == account_id:
                return account.get('daily_limit', 800)
        return 800  # 默认限额
    
    def mark_account_success(self, account_id: str, data_count: int):
        """标记账号使用成功"""
        if account_id not in self.account_status:
            logger.warning(f"[ACCOUNT_POOL] 账号状态不存在: {account_id}")
            return
        
        status = self.account_status[account_id]
        stats = self.usage_stats[account_id]
        
        # 更新状态
        status['status'] = 'available'
        status['usage_count_today'] += data_count
        status['total_data_crawled'] += data_count
        status['consecutive_errors'] = 0  # 重置错误计数
        status['last_success_time'] = time.time()
        
        # 更新统计
        stats['total_sessions'] += 1
        stats['successful_sessions'] += 1
        stats['total_data'] += data_count
        stats['success_rate'] = stats['successful_sessions'] / stats['total_sessions']
        
        logger.info(f"[ACCOUNT_POOL] ✅ {account_id} 成功: +{data_count} 条，今日总计: {status['usage_count_today']}")
    
    def mark_account_error(self, account_id: str, error_type: str):
        """标记账号使用失败"""
        if account_id not in self.account_status:
            logger.warning(f"[ACCOUNT_POOL] 账号状态不存在: {account_id}")
            return
        
        status = self.account_status[account_id]
        stats = self.usage_stats[account_id]
        current_time = time.time()
        
        # 增加错误计数
        status['consecutive_errors'] += 1
        stats['total_sessions'] += 1
        
        # 根据错误类型设置冷却
        if error_type == '403':
            # 403错误：指数退避冷却
            cooling_time = self.security_config.get('account_cooling_time', 3600)
            multiplier = self.security_config.get('error_cooling_multiplier', 2)
            actual_cooling = cooling_time * (multiplier ** (status['consecutive_errors'] - 1))
            
            status['status'] = 'cooling'
            status['cooling_until'] = current_time + actual_cooling
            
            logger.warning(f"[ACCOUNT_POOL] ⚠️ {account_id} 403错误，冷却 {actual_cooling/3600:.1f} 小时")
            
        elif error_type == 'captcha':
            # 验证码：固定冷却2小时
            cooling_time = 7200  # 2小时
            status['status'] = 'cooling'
            status['cooling_until'] = current_time + cooling_time
            
            logger.warning(f"[ACCOUNT_POOL] ⚠️ {account_id} 验证码，冷却 2 小时")
            
        elif error_type == 'network':
            # 网络错误：短时间冷却
            cooling_time = 1800  # 30分钟
            status['status'] = 'cooling'
            status['cooling_until'] = current_time + cooling_time
            
            logger.warning(f"[ACCOUNT_POOL] ⚠️ {account_id} 网络错误，冷却 30 分钟")
            
        else:
            # 其他错误：标准冷却
            cooling_time = self.security_config.get('account_cooling_time', 3600)
            status['status'] = 'cooling'
            status['cooling_until'] = current_time + cooling_time
            
            logger.warning(f"[ACCOUNT_POOL] ⚠️ {account_id} {error_type}错误，冷却 1 小时")
        
        # 更新成功率
        if stats['total_sessions'] > 0:
            stats['success_rate'] = stats['successful_sessions'] / stats['total_sessions']
    
    def get_pool_status(self) -> Dict:
        """获取账号池状态"""
        status = {
            'company_pool': {
                'total': len(self.network_pools['company']),
                'available': 0,
                'in_use': 0,
                'cooling': 0,
                'blocked': 0
            },
            'home_pool': {
                'total': len(self.network_pools['home']),
                'available': 0,
                'in_use': 0,
                'cooling': 0,
                'blocked': 0
            },
            'total_today_usage': 0,
            'total_data_crawled': 0
        }
        
        # 统计各网络池状态
        for network_type in ['company', 'home']:
            for account in self.network_pools[network_type]:
                account_id = account['id']
                account_status = self.account_status[account_id]['status']
                
                if self.is_account_available(account_id):
                    status[f'{network_type}_pool']['available'] += 1
                elif account_status == 'in_use':
                    status[f'{network_type}_pool']['in_use'] += 1
                elif account_status == 'cooling':
                    status[f'{network_type}_pool']['cooling'] += 1
                elif account_status == 'blocked':
                    status[f'{network_type}_pool']['blocked'] += 1
                
                # 累计统计
                status['total_today_usage'] += self.account_status[account_id]['usage_count_today']
                status['total_data_crawled'] += self.account_status[account_id]['total_data_crawled']
        
        return status
    
    def log_pool_status(self, network_type: str = None):
        """记录账号池状态"""
        if network_type:
            pools = [network_type]
        else:
            pools = ['company', 'home']
        
        for pool_type in pools:
            logger.info(f"[ACCOUNT_POOL] {pool_type.upper()} 账号池状态:")
            for account in self.network_pools[pool_type]:
                account_id = account['id']
                status = self.account_status[account_id]
                
                cooling_info = ""
                if status['status'] == 'cooling':
                    remaining = max(0, status['cooling_until'] - time.time())
                    cooling_info = f", 冷却剩余: {remaining/60:.1f}分钟"
                
                logger.info(f"  {account_id}: {status['status']}, "
                          f"今日: {status['usage_count_today']}, "
                          f"错误: {status['consecutive_errors']}{cooling_info}")
    
    def reset_daily_usage(self):
        """重置每日使用量（每日0点调用）"""
        logger.info("[ACCOUNT_POOL] 重置每日使用量统计")
        
        for account_id in self.account_status:
            self.account_status[account_id]['usage_count_today'] = 0
            
            # 如果账号状态正常，重置连续错误计数
            if self.account_status[account_id]['status'] == 'available':
                self.account_status[account_id]['consecutive_errors'] = 0
    
    def get_account_by_id(self, account_id: str) -> Optional[Dict]:
        """根据ID获取账号信息"""
        for account in self.accounts:
            if account['id'] == account_id:
                return account
        return None
    
    def get_network_type_by_account(self, account_id: str) -> Optional[str]:
        """根据账号ID获取网络类型"""
        account = self.get_account_by_id(account_id)
        if account:
            return account['network_binding']
        return None
